#coding=utf-8
# USAGE

import argparse
import cv2
import numpy as np
import time,sys
import imutils

SAVE_PHOTO = True
dir = sys.path[0]
faceCascade = cv2.CascadeClassifier(dir+'/models/cascade.xml')

class pokemon_tracker:
 
	def __init__(self,realtime,tb_indx = 0):
		self.DEBUG = False
		self.picHeight = 0
		self.picLength = 0
		self.conTourColor = (0,0,255) # red
		self.realTime = realtime # online time counting
		self.startTime= time.time()
		self.saved = False
		self.start = False # start counting time
		self.Found = False # find target now
		self.x = 0
		self.y = 0
		self.timeToSave = 1.8
		self.tb_indx = tb_indx
		self.saved_goal = []
		self.cur_goal = ''
		self.last_save_time = 0
	
	# check whether input contour could be a square 
	def checkSquare(self,w,h):
		if w >= h:
			Long = w
			Short = h
		else:
			Long = h
			Short = w
		rate = Short*1.0/Long
		if self.DEBUG:
			print('square rate:',rate)
		if rate < 0.7:
			return False,rate
		return True,rate

	def model_detect(self):
		frame = self.image.copy()
		gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
		rect = faceCascade.detectMultiScale(
			gray,
			scaleFactor=1.1,
			minNeighbors=5,
			minSize=(30,30)
			# flags = cv2.IMREAD_GRAYSCALE
		)
		[x,y,w,h] = [0,0,0,0]
		area = 0
		for (x_, y_, w_, h_) in rect:
			if w_*h_ > area:
				area = w_*h_
				[x,y,w,h] = [x_, y_, w_, h_]
		return x,y,w,h
			

	def findPotentialContour(self,input_image):
		image_, contours, hierarchy = cv2.findContours(input_image,1,2)
		judge_score = 0
		areaRate = 0
		w_h_rate = 0
		[x,y,w,h] = [0,0,0,0]
		for n in range(len(contours)):
			cnt = contours[n]
			area = cv2.contourArea(cnt)
			areaRate_ = area*1.0/(self.picHeight*self.picLength)

			if areaRate_ > 0.035 and areaRate_ < 0.06:
			
				x_,y_,w_,h_=cv2.boundingRect(cnt)
				if h_< (self.picHeight*1.0/6):
					continue
				

				# remove potential border noise
				if x_ < 10 or y_ < 10 or abs(x_+w_ - self.picLength) < 10 or abs(x_+w_ - self.picLength) < 10:
					continue

				flag,w_h_rate_ = self.checkSquare(w_,h_)
				if not flag:
					continue
				
				else:
					cur_judge_score = self.cal_score(w_h_rate_, area)
					if cur_judge_score > judge_score:
						judge_score = cur_judge_score
						[x,y,w,h] = [x_,y_,w_,h_]
						w_h_rate = w_h_rate_
						areaRate = areaRate_
		# print '%f %f' %(areaRate,w_h_rate)
		if areaRate < 0.04 and w_h_rate < 0.8: # bad catch. try cv classifier.
			# print 'bad detect'
			[x,y,w,h] = self.model_detect()
			judge_score = 1
			
		return judge_score,[x,y,w,h]

	def cal_score(self, w_h_rate, area):
		return (w_h_rate)**2*area
	# def savePic(self):
	# 	self.startTime= time.time()
	# 	self.conTourColor = (0,255,0) # green
	# 	name = str(self.time)+'.png'

	# check target in the middle of the screen
	def checkInside(self):
		if self.x > self.picLength/10 and self.y >= 4 and self.x < 9*self.picLength/10 and self.y <7*self.picHeight/8:
			return True
		# else:
		#  	print('self x,y %d,%d pic Len&Height:%d %d ' %(self.x,self.y,self.picLength,self.picHeight))
			#print('save start:',self.saved,self.start)
		return False

	def checkSaved(self):
		curTime = time.time()
		if self.checkInside(): # detect target inside red box
			
			#print(self.saved,' ',self.start)
			if not self.saved and not self.start:
				# if self.DEBUG:
				print('start timing',curTime)
				self.startTime= curTime
				self.start = True
			if curTime-self.startTime> self.timeToSave:
				self.conTourColor = (0,255,0) # green
				if self.Found:
					cv2.putText(self.image, 'SAVED', (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
				1, self.conTourColor, 2)
				self.last_save_time = curTime
			elif not self.saved and self.start:
				cv2.putText(self.image, 'detecting', (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
				1, self.conTourColor, 2)
			
			
			# if  curTime-self.startTime> 1.8 and not self.saved:
			# 	name = str(self.time)+'.png'
			# 	print('figure saved as ',name)
			# 	cv2.imwrite(name, self.image, [int(cv2.IMWRITE_JPEG_QUALITY),95])
			# 	self.saved = True
		# remove flash case.
		elif self.saved  and curTime-self.last_save_time> self.timeToSave*1.5: # saved. now moving out of border
			# if self.DEBUG:
			print('timing end')
			self.saved = False
			self.start = False
			self.conTourColor = (0,0,255) # red

		elif self.start and curTime-self.startTime> self.timeToSave*1.5: # no detection. cancel
			self.saved = False
			self.start = False
			self.conTourColor = (0,0,255) # red
			# if self.DEBUG:
			print('timing end. no detection')
		# print curTime - self.startTime

	def plotRectangle(self,x,y,w,h):
		
		if self.Found and self.checkInside():
			# rectangle around pokemon
			cv2.rectangle(self.image,(x+5,y+5),(x+w-5,y+h-5),self.conTourColor,4) # 画框
			# red rectangle outsite picture
			cv2.rectangle(self.image,(x-15,y-15),(x+w+15,y+h+15),(0,0,255),4) 	
		else:
			# static red rectangle
			cv2.rectangle(self.image,(self.picLength/6,self.picHeight/6),(5*self.picLength/6,5*self.picHeight/6),(0,0,255),4) 	

	# plot rectangle around the detected picture
	def findPokemon(self,image,vel_static): 
		# self.image = image
		self.image = imutils.resize(image, width=500)
		self.picHeight = self.image .shape[0]
		self.picLength = self.image .shape[1]

		if vel_static == False or self.cur_goal == '':
			cv2.rectangle(self.image,(self.picLength/6,self.picHeight/6),(5*self.picLength/6,5*self.picHeight/6),(0,0,255),4) 	
			# print 'moving'
			return self.image
		
		if self.DEBUG:
			print('image size:',self.picHeight ,' ',self.picLength)
		if self.realTime:
			self.checkSaved()
		
		gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
		avg_gray = np.sum(gray)/(gray.shape[0]*gray.shape[1])
		if self.DEBUG:
			print('avg gray:'+str(avg_gray))

		# modify threshold paremeters according to gray average scale.
		threshold = int(avg_gray)+10
		if avg_gray < 90:
			threshold = 88
		elif avg_gray > 100 and avg_gray < 125:
			threshold = 125
		elif avg_gray > 125:
			threshold = 125
		if self.DEBUG:
			print('threshold:',threshold)


		kernel_size=3
		blurred = cv2.GaussianBlur(gray, (kernel_size, kernel_size), 0)
		thresh = cv2.threshold(blurred, threshold, 255, cv2.THRESH_BINARY)[1] 
		if self.DEBUG:
			cv2.imshow("thresh", thresh)

		# closing process
		close_ker_num = 17 # 17 works for most cases
		kernel = np.ones((close_ker_num,close_ker_num),np.uint8) 
		closing = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)

		if self.DEBUG:
			cv2.imshow("closing", closing)

		judge_score,[x,y,w,h] = self.findPotentialContour(closing)

		# plot rectangle according to return contour

		
		if judge_score == 0:
			 # check for brick wall using thresh instead of closing
			if self.DEBUG:
				print('reCheck\n')
			judge_score,[x,y,w,h] = self.findPotentialContour(thresh)

		if judge_score == 0:
			self.Found = False
			[self.x,self.y] = [0,0]
		else:
			self.Found = True
			[self.x,self.y] = [x,y]

		self.plotRectangle(x,y,w,h)
		if self.checkInside() and not self.saved and self.start and time.time()-self.startTime> 1.2*self.timeToSave:
			name = self.cur_goal+'.png'
			if name not in self.saved_goal:
				print 'figure saved as ',name
				self.saved_goal.append(name)
				if SAVE_PHOTO:
					cv2.imwrite(name, self.image, [int(cv2.IMWRITE_JPEG_QUALITY),95])
				self.saved = True
				self.start = False
	
			
		return self.image


